package matroska

import (
	"fmt"
	"io"
	"log"
	"math"
	"time"

	"gitee.com/general252/ebml-go/mkv"
	"gitee.com/general252/ebml-go/mkvcore"
	"gitee.com/general252/ebml-go/webm"
)

// https://www.matroska.org/technical/elements.html
// https://www.webmproject.org/docs/container/

type WriteSeekCloser interface {
	io.Writer
	io.Seeker
	io.Closer
}

type Matroska struct {
	writerSu *writerFileSize
	writer   io.WriteCloser

	isMatroska bool

	tracks []Track
	err    error
}

func Open(w WriteSeekCloser, tracks ...Track) (*Matroska, error) {
	if len(tracks) == 0 {
		return nil, ErrNoTracks
	}

	writer := &writerFileSize{
		w: w,
	}

	m := &Matroska{
		writerSu:   writer,
		writer:     writer,
		tracks:     tracks,
		isMatroska: true,
	}

	m.modifyTrackNumber()

	if err := m.open(); err != nil {
		return nil, err
	}

	return m, nil
}

func OpenWebM(writer io.WriteCloser, tracks ...Track) (*Matroska, error) {
	if len(tracks) == 0 {
		return nil, ErrNoTracks
	}

	m := &Matroska{
		writer:     writer,
		tracks:     tracks,
		isMatroska: false,
	}

	m.modifyTrackNumber()

	if err := m.open(); err != nil {
		return nil, err
	}

	return m, nil
}

func (m *Matroska) open() error {

	var (
		w         = m.writer
		wemTracks []*webm.TrackEntry
		opts      = m.getOptions()
	)

	for _, track := range m.tracks {
		t := track.GetTrackEntry()
		wemTracks = append(wemTracks, t)
	}

	writers, err := webm.NewSimpleBlockWriter(w, wemTracks, opts...)
	if err != nil {
		return err
	}

	if len(writers) == 0 || len(writers) != len(m.tracks) {
		return fmt.Errorf("open error. error writers  %v %v", len(writers), len(m.tracks))
	}

	for i := 0; i < len(m.tracks); i++ {
		m.tracks[i].setBlockWriteCloser(writers[i])
	}

	return nil
}

func (m *Matroska) Close() {
	for _, track := range m.tracks {
		_ = track.GetWriter().Close()
	}

	if w := m.writer; w != nil {
		_ = w.Close()
	}
}

func (m *Matroska) FileSize() int {
	if w := m.writerSu; w != nil {
		return w.FileSize()
	}
	return -1
}

func (m *Matroska) GetTracks() []Track {
	return m.tracks
}

func (m *Matroska) WriteTrack(t Track, timestamp time.Duration, b []byte, keyframe ...bool) (int, error) {
	return t.Write(timestamp, b, keyframe...)
}

func (m *Matroska) WriteVideo(timestamp time.Duration, b []byte) (int, error) {
	if m.err != nil {
		return -1, m.err
	}

	for _, track := range m.tracks {
		if track.IsVideo() {
			return track.Write(timestamp, b)
		}
	}

	return -1, ErrNotFoundTrack
}

func (m *Matroska) WriteAudio(timestamp time.Duration, b []byte) (int, error) {
	if m.err != nil {
		return -1, m.err
	}

	for _, track := range m.tracks {
		if track.IsAudio() {
			return track.Write(timestamp, b)
		}
	}

	return -1, ErrNotFoundTrack
}

func (m *Matroska) WriteSubtitle(timestamp time.Duration, duration time.Duration, b []byte) (int, error) {
	if m.err != nil {
		return -1, m.err
	}

	for _, track := range m.tracks {
		if track.IsSubtitle() {
			return track.WriteWithDuration(timestamp, duration, b)
		}
	}

	return -1, ErrNotFoundTrack
}

func (m *Matroska) getOptions() []mkvcore.BlockWriterOption {
	var opts = []mkvcore.BlockWriterOption{
		mkvcore.WithCues(true),
		mkvcore.WithEBMLHeader(mkv.DefaultEBMLHeader),
		mkvcore.WithSegmentInfo(mkv.Info{
			TimecodeScale: 1000000, // 1ms
			MuxingApp:     "ebml-go.mkv.BlockWriter",
			WritingApp:    "ebml-go.mkv.BlockWriter",
			Duration:      0,
			DateUTC:       time.Now(),
		}),
		mkvcore.WithOnErrorHandler(func(err error) {
			log.Println(err)
			if err != nil && m.err == nil {
				m.err = err
			}
		}),
		mkvcore.WithOnFatalHandler(func(err error) {
			log.Println(err)
			if err != nil && m.err == nil {
				m.err = err
			}
		}),
	}

	if m.isMatroska {
		opts = append(opts, mkvcore.WithSeekHead(true))
		// opts = append(opts, mkvcore.WithSegmentInfo(mkv.DefaultSegmentInfo))
		// opts = append(opts, mkvcore.WithMarshalOptions(ebml.WithDataSizeLen(2), ebml.WithElementWriteHooks(func(element *ebml.Element) {})))
		//mkvcore.WithMarshalOptions(ebml.WithDataSizeLen(2)),
	}

	var (
		hasVideo = false
		hasAudio = false
	)
	for _, t := range m.tracks {
		if t.IsAudio() {
			hasAudio = true
		} else if t.IsVideo() {
			hasVideo = true
		}
	}
	if hasVideo {
		opts = append(opts, mkvcore.WithMaxKeyframeInterval(1, 900*0x6FFF))
	} else if hasAudio {
		opts = append(opts, mkvcore.WithMaxKeyframeInterval(1, math.MaxInt16-5000))
	}

	opts = append(opts, mkvcore.WithSimpleTags(
		mkvcore.Tag{
			Targets: mkvcore.Targets{},
			SimpleTag: []mkvcore.SimpleTag{
				{
					TagName:   "company",
					TagString: "xxx",
				},
				{
					TagName:   "format",
					TagString: fmt.Sprintf("Matroska: %v, WebM: %v", m.isMatroska, !m.isMatroska),
				},
			},
		},
	))

	return opts
}

func (m *Matroska) modifyTrackNumber() {
	var (
		videoTrackList    []Track
		audioTrackList    []Track
		subtitleTrackList []Track
		otherTrackList    []Track
	)

	for _, t := range m.tracks {
		track := t
		if t.IsVideo() {
			videoTrackList = append(videoTrackList, track)
		} else if t.IsAudio() {
			audioTrackList = append(audioTrackList, track)
		} else if t.IsSubtitle() {
			subtitleTrackList = append(subtitleTrackList, track)
		} else {
			otherTrackList = append(otherTrackList, track)
		}
	}

	var trackNumber = 1
	for _, t := range videoTrackList {
		t.setTrackNumber(trackNumber)
		trackNumber++
	}

	for _, t := range audioTrackList {
		t.setTrackNumber(trackNumber)
		trackNumber++
	}

	for _, t := range subtitleTrackList {
		t.setTrackNumber(trackNumber)
		trackNumber++
	}

	for _, t := range otherTrackList {
		t.setTrackNumber(trackNumber)
		trackNumber++
	}

}
